import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  
import os
import pandas as pd
import output_function as of
from operatorset import compute_stream_xperiod,apply_bnd_omega_yfreeslip_xperiod
'''只模拟密度的扰动量'''

'''输入速度，计算涡度。采用中心差分方法。
函数本义是通用的，但此程序中由于没有更新u,v在中间区域的值，所以就在设初值时用了一次
注意，左侧指标是j，右侧指标是i
最上方：[-1, 1:-1]
'''
def vel_to_omega(u, v, dx, dy):  
    omega = (v[1:-1, 2:] - v[1:-1, :-2]) / 2 / dx - (u[2:, 1:-1] - u[:-2, 1:-1]) / 2 / dy  
    return omega




#根据流函数计算u,v
def psi_to_vel(psi,u,v,dx,dy):
    #内点
    u[1:-1,1:-1]=(psi[2:, 1:-1] - psi[:-2, 1:-1])/2/dy
    v[1:-1,1:-1]=-(psi[1:-1, 2:] - psi[1:-1, :-2])/2/dx
    #边界
    u[:,0]=u[:,-2]
    u[:,-1]=u[:,1]
    u[-1, 1:-1] = u[-2, 1:-1]
    u[0, 1:-1] = u[1, 1:-1]

    v[:,0]=v[:,-2]
    v[:,-1]=v[:,1]
    v[-1, 1:-1] = 0 
    v[0, 1:-1] = 0
    return u,v

#对流算符
def convection_operator(myfun,u,v):
    #非守恒形式：
    uw_x = u[1:-1,1:-1] / 2 / dx  * (myfun[1:-1, 2:] - myfun[1:-1, :-2])
    vw_y=v[1:-1,1:-1] / 2 / dy  * (myfun[2:,1:-1] - myfun[:-2,1:-1])

    #守恒形式
    # uw_x =  (u[1:-1, 2:]*myfun[1:-1, 2:] -u[1:-1, :-2]*myfun[1:-1, :-2])/ 2 / dx
    # vw_y =  (v[2:, 1:-1]*myfun[2:, 1:-1] - v[:-2, 1:-1]*myfun[:-2, 1:-1])/2 / dy

    return uw_x,vw_y

#扩散算符
def diff_operator(myfun,nu,dx,dy):
    diff_y = nu * (myfun[2:, 1:-1] + myfun[:-2, 1:-1] - 2 * myfun[1:-1, 1:-1]) / dy / dy
    diff_x = nu * (myfun[1:-1, 2:] + myfun[1:-1, :-2] - 2 * myfun[1:-1, 1:-1]) / dx / dx
    return diff_x,diff_y


#漂移驱动
def drive_term_omega(n_density,beta):
    linear_flag=True
    if linear_flag:
        drt=-beta*(n_density[1:-1,2:]-n_density[1:-1,:-2])/2/dx
        return drt
    else:
        drt2=np.zeros_like(n_density)
        j_up=-1 #控制浮力驱动项存在范围，但实际上应当包含整个求解域。
        n_log_density=np.log(n_density)
        #程序中的x方向差分，但确实是沿极向的
        n_log1=n_log_density[1:j_up,2:]
        n_log2=n_log_density[1:j_up,:-2]
        #drt2=beta*(n_log[2:,1:-1]-n_log[:-2,1:-1])/2/dy
        drt2[1:j_up,1:-1]=-beta*(n_log1-n_log2)/2/dy    
        return drt2[1:-1,1:-1]


#涡度向前演化
def forward_omega_FTCS(omega,n_density,u,v):
    uw_x,vw_y =convection_operator(omega,u,v)
    diff_x,diff_y=diff_operator(omega,prandtl,dx,dy)  #扩散项用普朗特数
    #omega_drive_term=drive_term_omega(n_density,Ra*prandtl)
    omega_drive_term=-Ra*prandtl*(n_density[1:-1,2:]-n_density[1:-1,:-2])/2/dx
    return omega[1:-1, 1:-1] + dt * (-uw_x  -vw_y+omega_drive_term + diff_y + diff_x)

#密度向前演化
def forward_n_FTCS(n_density,u,v):
    uw_x,vw_y =convection_operator(n_density,u,v)
    diff_x,diff_y=diff_operator(n_density,1,dx,dy)    #归一化的密度扩散系数是1
    return n_density[1:-1, 1:-1] + dt * (-uw_x -vw_y+ diff_y + diff_x+v[1:-1, 1:-1])

#边界涡度处理，采用无滑移条件，顶盖速度恒为1，其余三壁速度恒为零.(太耍赖了，用psi和u,v)
def apply_bnd_omega_(psi, u, v, omega):
    #左右边界周期条件
    omega[:,0]=omega[:,-2]
    omega[:,-1]=omega[:,1]

    #顶部底部涡量直接设为零
    omega[0, 1:-1] =0
    omega[-1, 1:-1] =0

def apply_bnd_n_density(n_density):
    #顶边第一类条件
    n_density[-1,:]=n_up
    #底边第一类条件
    n_density[0,:]=n_down

    #左右边界周期条件
    n_density[:,0]=n_density[:,-2]
    n_density[:,-1]=n_density[:,1]


def diag_n_y_sec(n_density,out_iter):
    
    subfolder='n_density_y_section' #子文件夹名
    full_path_dir = os.path.join(os.getcwd(), subfolder) # 拼接当前工作目录和子文件夹路径 
    # 检查子文件夹是否存在  
    if not os.path.exists(full_path_dir):  
        # 如果不存在，则创建子文件夹  
        os.makedirs(full_path_dir)

    t_diag=out_iter*dt
    plt.figure(figsize=(12, 8), dpi=100)  
    plt.plot(y,n_density[:,Nx//2])
    plt.xlabel('y')  
    plt.ylabel('n_density') 
    plt.title('Ra={},t={}'.format(Ra,t_diag))  
    plt.grid(True)
    # 定义文件名，包含idiag的值  
    filename_base = 'n_density_y_plot_step_'  
    file_extension = '.png'  
    filename = f"{filename_base}{out_iter}{file_extension}"  
  
    # 拼接完整的文件路径  
    full_path = os.path.join(os.getcwd(), subfolder, filename)  
  
    # 保存散点图到指定路径  
    plt.savefig(full_path)  
    plt.close()

def save_simulation_data(n, psi, omega, u, v, out_iter, base_path='./data/SOL/'):
    """
    将模拟数据保存到.npz文件。

    参数:
    n, psi, omega, u, v -- 模拟数据的NumPy数组。
    x, y -- 空间网格的NumPy数组。
    out_iter -- 当前的时间步。
    base_path -- 保存文件的基础路径。
    """
    # 构建保存路径，包含时间步信息
    save_path = f"{base_path}/simulation_data_t_{out_iter}.npz"
    
    # 确保目录存在
    directory = os.path.dirname(save_path)
    if not os.path.exists(directory):
        os.makedirs(directory)
    
    # 保存数据
    np.savez_compressed(save_path, n=n, psi=psi, omega=omega, u=u, v=v,out_iter=out_iter)

base_folder=base_path = './data/RB_perturb/'  # 根据需要修改为您的保存路径
config_filename = f"{base_path}simulation_config.json"
config = of.load_simulation_config(config_filename)
# 从config字典中提取特定的参数并赋值给变量
dx = config['dx']
dy = config['dy']
Lx = config['Lx']
Ly = config['Ly']
Nx = config['Nx']
Ny = config['Ny']
n_up = config['n_up']
n_0 = config['n_0']
n_down = config['n_down']
Delta_n = config['Delta_n']
Ra=Ra_star = config['Ra_star']
prandtl=Pr = config['Pr']
dt = config['dt']
ntime = config['ntime']
ndiag = config['ndiag']
ndiag=50
# 使用示例
out_iter = 6000  # 假设我们要加载第300时间步的数据

# 加载数据
n, psi, omega, u, v, out_iter_loaded = of.load_simulation_data(base_path, out_iter)
n_init0_matrix, xx, yy,x,y = of.load_initial_data(load_path=f'{base_path}initial_conditions.npz')
n_density=n



if __name__ == '__main__':

    ntime=16000
    #进入主循环
    end_signal = False  #终止判断条件。可以再加一个判断条件。
    #while end_signal==False and out_iter<1000:
    while  out_iter<ntime:
        if out_iter % ndiag==0:
            #diag_n_y_sec(n_density,out_iter)    #截面密度分布
            #of.diag_n_pcolor(n_density,n_init0_matrix,u,v,xx,yy,out_iter,Ra,dt)
            """保存n,psi,omega,u,v到.npz"""
            save_simulation_data(n_density, psi, omega, u, v, out_iter,base_folder)
            print(f'out_iter:{out_iter}')

        out_iter+=1
        omega[1:-1, 1:-1] = forward_omega_FTCS(omega,n_density,u,v)  #时间向前步进。此时边界涡量会向中心扩散
        n_density[1:-1, 1:-1]=forward_n_FTCS(n_density,u,v)
        psi, end_signal = compute_stream_xperiod(omega[1:-1, 1:-1],Nx,Ny, dx, dy, psi_init=psi, tol=1e-5)  #泊松方程计算流函数。算出的流函数将被用于下一次边界设置和时间步进。
        #需要新增代码：根据流函数场计算速度场。
        u,v=psi_to_vel(psi,u,v,dx,dy)
        apply_bnd_omega_(psi, u, v, omega)  #设置边界涡量。u,v仅用到边界值，其不在程序中改变。因此这里的u,v并不是全流场速度。
        apply_bnd_n_density(n_density)

        
      